package com.googlecode.yoohoo.io.socket.stream;

import java.io.StringReader;

import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamConstants;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;

import org.apache.mina.filter.SSLFilter;

import com.googlecode.yoohoo.io.socket.security.TlsEnvironment;
import com.googlecode.yoohoo.xmppcore.protocol.XmppProtocolException;
import com.googlecode.yoohoo.xmppcore.protocol.stream.error.NotAuthorized;
import com.googlecode.yoohoo.xmppcore.protocol.stream.error.XmlNotWellFormed;
import com.googlecode.yoohoo.xmppcore.stream.INegotiationContext;

public class TlsNegotiant extends InitialStreamNegotiant {
	private static final String NAME_SSL_FILTER = "SSLFilter";
	private static final String RESPONSE_TLS_PROCEED = "<proceed xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>";
	private static final String XMPP_STARTTLS_TAG_NAME = "starttls";
	private static final String XMPP_TLS_NAMESPACE = "urn:ietf:params:xml:ns:xmpp-tls";
	private static SSLFilter sslFilter;
	private static Object sslContextInitLock = new Object();
	
	private String appHome;
	private String hostName;
	private boolean tlsRequired;

	enum TlsNegotiationState {
		NONE,
		PROCEED_RESPONSE_SENT,
		TLS_ESTABLISHED
	};
	
	private TlsNegotiationState state;
	
	public TlsNegotiant(IStreamFeaturesManager featuresManager, String appHome,
			String hostName, boolean tlsRequired) {
		super(featuresManager);
		this.appHome = appHome;
		this.hostName = hostName;
		this.tlsRequired = tlsRequired;
		state = TlsNegotiationState.NONE;
	}
	
	@Override
	protected void doNegotiate(INegotiationContext context, String message) {
		if (state != TlsNegotiationState.TLS_ESTABLISHED && sslFilter != null &&
				sslFilter.isSSLStarted(((NegotiationContext)context).getIoSession())) {
			state = TlsNegotiationState.TLS_ESTABLISHED;
		}
		
		if (state == TlsNegotiationState.NONE) {
			processStartTlsRequest(context, message);
		} else if (state == TlsNegotiationState.TLS_ESTABLISHED) {
			super.doNegotiate(context, message);
		}
	}

	private void processStartTlsRequest(INegotiationContext context, String message) {
		if (isStartTlsRequest(message)) {
			try {
				synchronized (sslContextInitLock) {
					if (sslFilter == null) {
						TlsEnvironment tlsEnvironment = new TlsEnvironment(appHome, hostName);
						sslFilter = new SSLFilter(tlsEnvironment.getSslContext());
						sslFilter.setEnabledCipherSuites(new String[] {
								"TLS_DHE_RSA_WITH_AES_256_CBC_SHA",
								"TLS_DHE_DSS_WITH_AES_256_CBC_SHA",
								"TLS_RSA_WITH_AES_256_CBC_SHA",
								"TLS_DHE_RSA_WITH_AES_128_CBC_SHA",
								"TLS_DHE_DSS_WITH_AES_128_CBC_SHA",
								"TLS_RSA_WITH_AES_128_CBC_SHA"
						});
					}
				}
			} catch (SecurityException e) {
				context.write("<failure xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>");
				context.close();
				
				return;
			}
			
			((NegotiationContext)context).getIoSession().getFilterChain().addFirst(NAME_SSL_FILTER, sslFilter);
			((NegotiationContext)context).getIoSession().setAttribute(SSLFilter.DISABLE_ENCRYPTION_ONCE, true);			
			context.write(RESPONSE_TLS_PROCEED);
			
			state = TlsNegotiationState.PROCEED_RESPONSE_SENT;
		} else {
			if (tlsRequired) {
				throw new XmppProtocolException(new NotAuthorized());
			} else {
				nextPhase(context);
				if (nextNegotiant != null) {
					nextNegotiant.negotiate(context, message);
				}
			}
		}
	}

	private boolean isStartTlsRequest(String message) {
		XMLInputFactory factory = XMLInputFactory.newInstance();
		try {
			XMLStreamReader reader = factory.createXMLStreamReader(new StringReader(message));
			int eventType = reader.next();
			
			if (eventType != XMLStreamConstants.START_ELEMENT) {
				return false;
			}
			
			if (!XMPP_STARTTLS_TAG_NAME.equals(reader.getLocalName()))
				return false;
			
			if (reader.getNamespaceCount() != 1)
				return false;
			
			if (!XMPP_TLS_NAMESPACE.equals(reader.getNamespaceURI(0)))
				return false;
			
			eventType = reader.next();
			if (eventType != XMLStreamConstants.END_ELEMENT)
				return false;
			
			return true;
		} catch(XMLStreamException e) {
			throw new XmppProtocolException(new XmlNotWellFormed());
		}
	}

	@Override
	protected NegotiationPhase getPhase() {
		return NegotiationPhase.TLS;
	}
}
